package com.seguridad.txwrapping.tester;

import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URL;
import java.security.MessageDigest;

import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.soap.MessageFactory;
import javax.xml.soap.SOAPConnection;
import javax.xml.soap.SOAPConnectionFactory;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPMessage;
import javax.xml.transform.Result;
import javax.xml.transform.Source;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpression;
import javax.xml.xpath.XPathFactory;

import org.apache.log4j.Logger;
import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import com.seguridad.txwrapping.TestResult;

public abstract class WSSWTester {

	static Logger logger = Logger.getLogger(WSSWTester.class);
	private URL endpoint;
	private String validMessagePath;
	private String idToWrap;
	
	public WSSWTester(URL endpoint, String validMessagePath, String idToWrap) {
		this.endpoint = endpoint;
		this.validMessagePath = validMessagePath;
		this.idToWrap = idToWrap;
	}

	public TestResult test() {
		String okHash = sampleOK();
		String failHash = sampleIntegrityFail();
		String wrappedHash = sampleWrapped();
		
		String status;
		String message;
		
		if(wrappedHash.equals(okHash)){
			status = "VULNERABLE";
			message = "Se obtuvo la misma respuesta que un mensaje valido";
		} else if(!wrappedHash.equals(failHash)){
			status = "VULNERABLE";
			message = "Se obtuvo diferente respuesta que mensaje valido pero pasa integridad";
		} else {
			status = "NOT_VULNERABLE";
			message = "Se obtuvo la misma respuesta que un mensaje invalido por integridad";
		}
		
		return new TestResult(status, message);
	}


	private String sampleOK() {
		logger.debug("########## SAMPLE MENSAJE OK");
		String hash = null;
		InputStream is = null;
		SOAPConnection connection = null;
		try {
			is = new FileInputStream(getValidMessagePath());
			SOAPConnectionFactory sfc = SOAPConnectionFactory.newInstance();
			connection = sfc.createConnection();
				  
			MessageFactory mf = MessageFactory.newInstance();
			SOAPMessage sm = mf.createMessage(null, is);
			
			if(logger.isDebugEnabled()){
				logger.debug("REQUEST");
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				sm.writeTo(baos);
				logger.debug(baos.toString());
			}
			
			SOAPMessage response = connection.call(sm, getEndpoint());
			
			if(logger.isDebugEnabled()){
				logger.debug("RESPONSE");
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				response.writeTo(baos);
				logger.debug(baos.toString());
			}
			
			hash = hashBody(response);
			
		} catch (Exception ex) {
			ex.printStackTrace();
		} finally {
			try {
				if(connection != null) {connection.close();}
				if(is != null) {is.close();}
			} catch (SOAPException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}
			
		}
		return hash;
	}

	
	private String sampleIntegrityFail() {
		logger.debug("########## SAMPLE ERROR INTEGRIDAD");
		String hash = null;
		SOAPConnection connection = null;
		InputStream is = null;
		try {
			is = new FileInputStream(getValidMessagePath());
			SOAPConnectionFactory sfc = SOAPConnectionFactory.newInstance();
			connection = sfc.createConnection();
			
			MessageFactory mf = MessageFactory.newInstance();
			SOAPMessage sm = mf.createMessage(null, is);
			
			
			// Busca y modifica el hash de la referencia
			NodeList nl = sm.getSOAPPart().getDocumentElement().getElementsByTagNameNS(XMLSignature.XMLNS, "Reference");
			for (int i = 0; i < nl.getLength(); i++) {
				Node node = nl.item(i);
				String uri = node.getAttributes().getNamedItem("URI").getNodeValue();
				
				if(uri.equals(getIdToWrap()) || uri.equals("#"+getIdToWrap())){
					NodeList childs = node.getChildNodes();
					for (int j = 0; j < childs.getLength(); j++) {
						Node child = childs.item(j);
						if("DigestValue".equalsIgnoreCase(child.getLocalName())){
							child.setTextContent("111111133333335a9acb1f677844444442222222");						
						}	
					}
				}
			}
			
			if(logger.isDebugEnabled()){
				logger.debug("REQUEST");
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				sm.writeTo(baos);
				logger.debug(baos.toString());
			}
			
			SOAPMessage response = connection.call(sm, getEndpoint());
			
			if(logger.isDebugEnabled()){
				logger.debug("RESPONSE");
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				response.writeTo(baos);
				logger.debug(baos.toString());
			}
	
			hash = hashBody(response);
			
		} catch (Exception ex) {
			ex.printStackTrace();
		} finally {
			try {
				if(connection != null) {connection.close();}
				if(is != null) {is.close();}
			} catch (SOAPException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		return hash;
	}
	
	protected abstract String sampleWrapped();
	
	protected Node searchNodeToWrap(SOAPMessage sm) {
		Node theNode = null;
		try {
			// Busca en el documento un elemento con el id de la referencia
			XPath xpath = XPathFactory.newInstance().newXPath();
			XPathExpression xPathExpression = xpath.compile("//*[@Id='" + getIdToWrap() + "']");  
			Object resultObject = xPathExpression.evaluate(sm.getSOAPPart().getDocumentElement(), XPathConstants.NODESET);
			NodeList nodes = (NodeList) resultObject;
			theNode = nodes.item(0);
			
		} catch (Exception e) {
			e.printStackTrace();
		}
		return theNode;
	}
	
	protected String hashBody(SOAPMessage response) {
		String hash = null;
		try {
			Document doc = response.getSOAPBody().extractContentAsDocument();
			
			Transformer transformer = TransformerFactory.newInstance().newTransformer();
		    Source source = new DOMSource(doc);
		    OutputStream os = new ByteArrayOutputStream();
		    Result output = new StreamResult(os);
		    transformer.transform(source, output);
		    	    
		    String normalized = os.toString().trim().replaceAll(">\\s+<", "><").replaceAll("^<\\?xml[^?]+\\?>", "");
		    MessageDigest sha1 = MessageDigest.getInstance("SHA1");
		    byte[] digest = sha1.digest(normalized.getBytes());
	    
		    StringBuffer buf = new StringBuffer();
		    for (int i = 0; i < digest.length; i++) { 
		        int halfbyte = (digest[i] >>> 4) & 0x0F;
		        int two_halfs = 0;
		        do { 
		            if ((0 <= halfbyte) && (halfbyte <= 9)) 
		                buf.append((char) ('0' + halfbyte));
		        else 
		            buf.append((char) ('a' + (halfbyte - 10)));
		            halfbyte = digest[i] & 0x0F;
		        } while(two_halfs++ < 1);
		    }  
			
		    hash = buf.toString();
			
		} catch (Exception e) {
			e.printStackTrace();
		}	    
	   
		return hash;
	}
	
	
	public URL getEndpoint() {
		return endpoint;
	}

	public String getValidMessagePath() {
		return validMessagePath;
	}

	public String getIdToWrap() {
		return idToWrap;
	}

}
